function   results = find_solution_inner_iteration_test( p_2D,  par  )
% Usage: find_solution_inner_iteration( alpha,  beta,  p_2D,  par  ). 
% 
% Notes: p_2D is the current price function, expressed in a 2D matrix with both w and lambda dimensions.
% This function iterates from an earlier price matrix to a new price matrix.

%% SECTION 1. PARAMETERS
N_w = par.N_w;    N_lambda = par.N_lambda;
AH=par.AH;   AL=par.AL;   rho=par.rho;  sigmaK = par.sigmaK;   investment=par.investment;   
alpha=par.alpha;   beta=par.beta;
max_jump=par.max_jump;   
p0 = par.p0;  p1=par.p1;   
w_grid = par.w_grid;   lambda_grid = par.lambda_grid;
print_level = par.print_level;   % 0: nothing.  1: important outputs.  2: print each loop.  3: second inner loop.   4: third inner loop.

%% SECTION 2. One iteration on the whole 2D function
% Last round price function p_fitted
p_2D_fitted =  scatteredInterpolant(   kron(ones(N_lambda,1), w_grid ),    kron(lambda_grid, ones(N_w,1)) ,  reshape( p_2D, N_w*N_lambda, 1 ),   'linear');
par.p_2D_fitted = p_2D_fitted;

% Other output variables
init_matrix = zeros(  N_w,  N_lambda );
p_matrix = init_matrix;     kappap_matrix = init_matrix;  

% error control
epsilon=0.001; % maximum error in solving kappap

for( counter_lambda = 1:N_lambda)
    lambda = lambda_grid(counter_lambda);   par.lambda=lambda;  
%     disp( [   'Inside 2D iteration: lambda=', num2str(lambda)   ]   )
    
    kappap_sol = zeros( N_w, 1 ); xK_sol = zeros(N_w,1);  yK_sol = zeros(N_w,1);    psi_new = zeros(N_w,1);   solved_vec = zeros(N_w,1);
    p_new = p_2D(:,counter_lambda) ;      pw_new=derivative( w_grid, p_new );
    %== Step 1: For each grid point of w, solve kappa, xK, yK. 
    for(k = 2:(N_w-1))
        % Decide whether at p, and w we have psi=1
        w=w_grid(k);    p=p_new(k);   pw=pw_new(k);

        if(print_level>=2)
            disp( [ '** w=', num2str(w),  ',  k=', num2str(k) ] );
        end
        yK_fun = @(xK) ( 1 - w*xK ) ./ (1-w);

        sigmap_fun = @(xK) pw.*w.*(1-w).*(xK- yK_fun(xK) ) ./ ( 1- pw.*w.*(1-w).*(xK- yK_fun(xK) )  ).* sigmaK; 
        sigmah_fun = @(xK) yK_fun(xK) .* (sigmaK + sigmap_fun(xK) ); 
        sigma_fun = @(xK)  xK .* ( sigmaK + sigmap_fun(xK) );

        par.yK_fun=yK_fun;  par.sigmap_fun=sigmap_fun;  par.sigmah_fun=sigmah_fun;  par.sigma_fun=sigma_fun;
        par.w_ind = w;    par.p_ind = p;   par.pw_ind = pw;   

        xK_min = 1  ;     % The value of xK such that yK=xK
        if(pw<=0)
            xK_max1 =  1000 ;   % Don't limit the value of xK, and use a uniform set of first order conditions.
        else
            xK_max1 = 1/( pw*w) + 1 ;    % maximum value restricted by positive volatility
        end
        xK_max1 = min( xK_max1,  (1+alpha*beta)/(alpha*beta) ) ;  % Maximum likelihood restricted by kappa equation

        xKs = exp(linspace( log(xK_min+1), log(xK_max1*0.98+1) ,100)') - 1;
        if(print_level>=3  &&  w>0.1 )
            figure;
            subplot(2,2,1);plot( xKs, yK_fun(xKs) );  xlabel( 'x^K' ); title('y^K');  
            subplot(2,2,2);plot( xKs, sigmap_fun(xKs) );  xlabel( 'x^K' ); title('\sigma^p');  
            subplot(2,2,3);plot( xKs, sigmah_fun(xKs) );  xlabel( 'x^K' ); title('\sigma^h');   
            subplot(2,2,4);plot( xKs, sigma_fun(xKs) );  xlabel( 'x^K' ); title('\sigma'); 
            pause();
        end

        N_sample = 15;
        % kappaps = [0;  exp(linspace( log(0.001), log(0.9), N_sample-1 )')] ; 
        kappaps = linspace( 0, min( max_jump, p-0.02 ) , N_sample )' ; 
        xK_roots = zeros( numel(kappaps), 1 );
        solved_inner = (ones( N_sample,1 )==1);   
        for(i = 1:N_sample)
            kappap = kappaps(i);
            xK_max = min( xK_max1,   (1+alpha*beta)/(kappap+alpha*beta)  );    par.xK_max = xK_max;   % Note: we should use min(xK_max, 1/kappap) so that it updates for each kappap, instead of min(par.xK_max, 1/kappap)
            par.xK_min = xK_min;  
            [ xK_sol , fval ,  ~,  ~] = fsolve( @(xK) 10*xK_iterative_residual(xK, kappap, par) , xK_min*1.1 ,  optimset('Display','off') );  % note: scaling by 10 improves solution accuracy                
            if( abs(fval)>epsilon  )
                xK_corner = 1/w;
                if( xK_iterative_residual(1, kappap, par) * xK_iterative_residual(1.001, kappap, par) < 0 )
                    xK_sol = 1; solved_inner(i)=true;
                elseif( xK_iterative_residual(xK_corner, kappap, par) > 0  &&  xK_corner<=xK_max )   % Corner case works
                    xK_sol = xK_corner; solved_inner(i)=true;
                else   % Corner case does not work
                    disp( [ 'In solving xK1, a solution fails (should not happen)!!! At kappap=', num2str(kappap), '(iteration=', num2str(i)  , '), w=', num2str(w),  '(main iteration=', num2str(k), ')' ] );
                    xK_sol = xK_min;   solved_inner(i)=false;
                end
            end
            xK_roots(i)  =  xK_sol;
            if(print_level>=4 )    
                xKs = linspace( xK_min,  xK_max*0.99 , 100 )';   % At maximum, we should reach 1/kappap, which results in an infinitely negative residual                                
                residuals = xK_iterative_residual(xKs, kappap, par);  % Note: we don't have to change par.xK_min = xK_min, because it only controls the bundary and we are in the boundary    
                figure; plot( xKs  ,  residuals , xKs, 0*xKs, '-.',   xK_sol,   0, '*'  );  title( ['kappap=', num2str(kappaps(i))] );  xlabel('xK');  ylabel('residual');
                pause();
            end
        end
        close all;

        % roots_SLM = slmengine( kappaps(solved) ,xK_roots(solved), 'plot', 'on', 'decreasing', 'on',    'knots',   floor( N_sample /2)   );  
        roots_xK = interp1( kappaps(solved_inner) ,xK_roots(solved_inner),'linear','pp');

        % Then add parameters to par for solving kappap
        par.xK_roots_fitted = roots_xK;

        counter=0;   
        kp0=0;       max_kp =  max(kappaps)-0.01   ;  step_size=0.1;    maxit = floor(max_kp/step_size) ;  fval=1;
        while( abs(fval)>epsilon & counter<maxit & kp0<max_kp )   % exitflat<0 means no solution has been found. If so, we increase the initial value kp0 and search again.
                counter = counter+1;
                [kappap_sol(k), fval ,~,~] = fsolve( @(kp) kappa_iterative_residual(kp, par),  kp0, optimset( 'Display','none' ) );  % Important: use subscript k
                kp0=kp0+step_size;  % Note: it is very important to progressively increase the initial point. If we plot the residual as a function of kappa, then we can observe there are multiple solutions. We pick the solution closest to zero by this algorithm. 
        end
        solved= (abs(fval)<epsilon) ;
        if( ~solved )  
            disp( [ 'In solving kappap, a solution fails! At w=', num2str(w),  '(iteration=', num2str(k), ')' ] )
        end
        solved_vec(k) = solved;
        xK_sol(k) = ppval( roots_xK,kappap_sol(k));  yK_sol(k) = yK_fun( xK_sol(k)  );
        psi_new(k) = w.*xK_sol(k) ; 
        if( (print_level>=2 & w>0.01)  )  %  Plot the residuals for each kappa
            kappas_plot = [ linspace(0, 0.9*min(kappap_sol),100)'  ; kappaps ];
            residuals = kappa_iterative_residual( kappas_plot, par );  
            subplot( 1,2,1 );   plot(  kappas_plot,  ppval(roots_xK, kappas_plot), 'o-'  , kappap_sol(k), xK_sol(k), '*'   ); title( [ 'w=', num2str(w), ', k=', num2str(k)] ) ; ylabel('x^K'); xlabel('\kappa^p');
            subplot( 1,2,2 );   plot( kappas_plot, residuals,   'o-',  kappap_sol(k),  0, '*' ); title( [ 'w=', num2str(w),  ', k=', num2str(k)] ) ;  xlabel('\kappa^p');   ylabel('residual' );   
            pause();
        end  
    end

    if(print_level>=1)
        plot( w_grid, kappap_sol );  title( '\kappa^p(w)' );  xlabel('w');  pause();  
    end

    %== Step 2: Get updated results for price function. Note that psi<=1 must be satisfied.
    psi_sol = min( psi_new, 1 ) ;        psi_sol( w_grid==1 ) = 1;    psi_sol( w_grid==0 ) = 0; 
    p_sol_next = zeros(  numel(w_grid) ,1 );  
    for(i = 1:numel(w_grid))
        p_sol_next(i) = fsolve(  @(p)  investment(p) + rho*p - AL - psi_sol(i)*( AH-AL ),  1 ,   optimoptions( 'fsolve','Display','none' ) );  
    end

    % We need to smooth the function around zero
    psi_new(N_w)=1;  solved_vec=(solved_vec==1 & (~isnan(solved_vec)) );  solved_vec(1)=true;  solved_vec(N_w)=true;
    max_index = min( find( psi_new>=1 ) )  ;  % The minimum index for psi greater than 1.
    index = 1:(max_index);   index=  index( solved_vec(1:max_index)  );

    if (print_level>=1)    print_option = 'on' ;     else print_option = 'off'  ;       end
    p_SLM = slmengine( w_grid(index) ,  p_sol_next(index)  ,'plot', print_option, 'increasing', 'on',  'concavedown', 'on', 'leftvalue',p0, 'rightvalue',  p_sol_next(max(index))  , 'knots',   floor(max(index)/3)   );  
    if( print_level>=1 )  pause();  end;    p
    % Update p_vec
    p_vec = [ slmeval( w_grid(1:max_index) ,p_SLM);   ones( N_w-max_index, 1 )*p1  ] ;   % The smoothed price function. The latter part of price should be equal to the ode price
    kappap_vec = interp1(w_grid(solved_vec) ,  kappap_sol(solved_vec),  w_grid);

    p_matrix(:,counter_lambda) = p_vec;     
    kappap_matrix(:,counter_lambda) = kappap_vec;  
end


%% SECTION 2. OUTPUT
results.p_matrix = p_matrix;
results.kappap_matrix = kappap_matrix;


